# -*- coding: utf-8 -*-

import os
import sys
import re
import json
import logging
import time
from typing import List, Dict, Any, Optional

# 设置环境变量
os.environ['CUDA_VISIBLE_DEVICES'] = '4'

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 导入swift相关的类
from swift.llm import PtEngine, RequestConfig, InferRequest
from generate_reply_prompt import GenerateReplyPrompt

class ModelHandler:
    """
    模型处理器，负责加载模型和生成回答
    """
    
    def __init__(self, config):
        """
        初始化模型处理器
        
        Args:
            config: 配置对象
        """
        self.config = config
        self.engine = None
        
        # 初始化提示词生成器
        self.prompt_generator = GenerateReplyPrompt("/mnt/ssd/jsj/patient/script/人格定义.json")
        
    def load_model(self, model_path: str) -> None:
        """
        加载模型
        
        Args:
            model_path: 模型路径
        """
        try:
            logger.info(f"开始加载模型: {model_path}")
            
            # 加载推理引擎
            self.engine = PtEngine(model_path, max_batch_size=8)
            
            logger.info("模型加载成功")
            
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            raise
            
    def _parse_personality_info(self, personality_info: Any) -> Dict[str, str]:
        """
        解析人格信息为字典格式
        
        Args:
            personality_info: 人格信息（可能是字符串、字典或None）
        
        Returns:
            解析后的人格信息字典
        """
        target_personality = {}
        
        try:
            if isinstance(personality_info, str):
                # 预处理：将所有换行符转换为逗号，以便统一处理
                processed_text = personality_info.replace('\n', ',').strip()
                
                # 尝试JSON解析
                if processed_text.startswith('{') and processed_text.endswith('}'):
                    # 处理单引号JSON格式
                    if '\'' in processed_text and '"' not in processed_text:
                        processed_text = processed_text.replace('\'', '"')
                    try:
                        target_personality = json.loads(processed_text)
                        logger.debug(f"JSON解析成功: {target_personality}")
                    except json.JSONDecodeError:
                        # JSON解析失败，尝试宽松格式解析
                        processed_text = processed_text[1:-1].strip()
                        target_personality = self._parse_loose_format_personality(processed_text)
                else:
                    # 直接尝试宽松格式解析
                    target_personality = self._parse_loose_format_personality(processed_text)
            elif isinstance(personality_info, dict):
                target_personality = personality_info
                logger.debug(f"人格信息已是字典格式: {target_personality}")
            elif personality_info is None:
                logger.debug("人格信息为None")
            
            # 如果解析结果为空，设置默认值
            if not target_personality:
                logger.debug("解析结果为空，设置默认人格信息")
                target_personality = {"性格": "未知", "情绪状态": "未知"}
            
            logger.debug(f"最终人格信息解析结果: {target_personality}")
        except Exception as e:
            logger.warning(f"解析人格信息失败: {str(e)}")
            # 即使失败，也设置一个默认的人格信息以避免后续错误
            target_personality = {"性格": "未知", "情绪状态": "未知"}
        
        return target_personality
    
    def _parse_loose_format_personality(self, processed_text: str) -> Dict[str, str]:
        """
        尝试宽松格式解析人格信息
        
        Args:
            processed_text: 处理后的文本
        
        Returns:
            解析后的人格信息字典
        """
        target_personality = {}
        
        # 分割键值对（支持分号和逗号）
        if ';' in processed_text:
            pairs = [p.strip() for p in processed_text.split(';') if p.strip()]
        else:
            pairs = [p.strip() for p in processed_text.split(',') if p.strip()]
        
        for pair in pairs:
            if ':' in pair or '：' in pair:
                try:
                    key, value = pair.split(':', 1) if ':' in pair else pair.split('：', 1)
                    key = key.strip().strip('"\'')
                    value = value.strip().strip('"\'')
                    # 检查是否是我们关心的人格维度
                    if key in ["性格", "情绪状态", "病史回忆能力", "医学理解力", "语言表达能力", "沟通意愿"]:
                        target_personality[key] = value
                        logger.debug(f"宽松解析到 {key}: {value}")
                except ValueError:
                    # 分割失败，跳过这个键值对
                    logger.debug(f"无法解析键值对: {pair}")
        
        return target_personality
    
    def _parse_patient_info(self, basic_info: Any) -> Dict[str, str]:
        """
        解析患者基本信息
        
        Args:
            basic_info: 患者基本信息（可能是字符串、字典或None）
        
        Returns:
            解析后的患者基本信息字典
        """
        patient_info_dict = {}
        
        try:
            if isinstance(basic_info, str):
                # 简单解析患者基本信息
                for line in basic_info.split('\n'):
                    if ':' in line or '：' in line:
                        key, value = re.split(r'[:：]\s*', line.strip(), maxsplit=1)
                        patient_info_dict[key] = value
            elif isinstance(basic_info, dict):
                patient_info_dict = basic_info
        except Exception as e:
            logger.warning(f"解析患者基本信息失败: {str(e)}")
        
        return patient_info_dict
    
    def _clean_response(self, base_response: str) -> str:
        """
        清理生成的回答，确保按照规定的格式处理
        
        Args:
            base_response: 原始生成的回答
        
        Returns:
            清理后的回答
        """
        if not base_response:
            return "（根据您的问题，我需要详细描述我的症状和感受。）"
        
        # 1. 移除所有可能的前缀（如"患者："、"丁xx："等）
        prefixes = ["患者：", "丁xx：", "病人：", "张xx："]
        for prefix in prefixes:
            if base_response.startswith(prefix):
                base_response = base_response[len(prefix):]
                logger.debug(f"移除了前缀'{prefix}'，剩余内容: {base_response[:50]}...")
                break
        
        # 2. 查找第一个中文字符的位置，移除前面的所有非中文字符
        first_valid_pos = None
        for i, char in enumerate(base_response):
            # 只将中文字符视为有效字符
            if 0x4e00 <= ord(char) <= 0x9fff:
                first_valid_pos = i
                break
        
        if first_valid_pos is not None:
            if first_valid_pos > 0:
                # 如果找到了有效的中文字符位置，并且不是第一个字符，则移除前缀
                base_response = base_response[first_valid_pos:]
                logger.debug(f"清理了奇怪的前缀，剩余内容: {base_response[:50]}...")
        else:
            # 如果没有找到有效的中文字符，则返回一个默认的提示信息
            logger.debug(f"未找到有效的中文字符，使用默认提示")
            return "（根据您的问题，我需要详细描述我的症状和感受。）"
        
        # 3. 移除可能包含的医生提问或对话历史内容
        # 查找可能的医生提问开始位置
        doctor_keywords = ["医生：", "用户：", "你："]
        for keyword in doctor_keywords:
            idx = base_response.find(keyword)
            if idx != -1:
                base_response = base_response[:idx]
                logger.debug(f"移除了后续的医生提问，剩余内容: {base_response[:50]}...")
                break
        
        # 4. 移除可能包含的结束场景描述
        end_scene_keywords = ["离开诊室", "张xx离开了诊室", "请记住，在下一轮对话中"]
        for keyword in end_scene_keywords:
            idx = base_response.find(keyword)
            if idx != -1:
                base_response = base_response[:idx]
                logger.debug(f"移除了结束场景描述，剩余内容: {base_response[:50]}...")
                break
        
        # 5. 移除可能包含的多余对话轮次
        dialog_ending_keywords = ["再见", "祝您健康", "祝您早日康复"]
        for keyword in dialog_ending_keywords:
            idx = base_response.find(keyword)
            if idx != -1:
                # 只保留第一个再见之前的内容
                base_response = base_response[:idx + len(keyword)]
                logger.debug(f"移除了多余对话轮次，剩余内容: {base_response[:50]}...")
                break
        
        # 6. 清理末尾的无关内容和空白字符
        base_response = base_response.strip()
        
        # 7. 如果清理后回答为空，使用默认回答
        if not base_response:
            logger.debug(f"清理后回答为空，使用默认提示")
            base_response = "（根据您的问题，我需要详细描述我的症状和感受。）"
        
        return base_response

    import re

    def strip_think_tags(self,text: str) -> str:
        """
        删除 <think>...</think> 部分，只保留其后的模型输出内容
        """
        # 用正则匹配并替换
        cleaned = re.sub(r"^.*?</think>", "", text, flags=re.DOTALL)
        return cleaned.strip()


    def batch_generate_responses(
        self,
        batch_data: List[Dict[str, Any]]
    ) -> List[str]:
        """
        批量生成回答
        
        Args:
            batch_data: 批量数据
        
        Returns:
            生成的回答列表
        """
        start_time = time.time()
        responses = []
        
        # 检查模型是否已加载
        if self.engine is None:
            logger.error("模型引擎未加载，无法生成回答")
            return ["（模型未初始化，无法生成回答）"] * len(batch_data)
        
        # 设置生成配置
        max_tokens = 512
        temperature = 0
        request_config = RequestConfig(max_tokens=max_tokens, temperature=temperature)
        
        # 准备推理请求
        infer_requests = []
        
        try:
            logger.info(f"准备批量请求，总数量: {len(batch_data)}")
            
            for i, data in enumerate(batch_data):
                logger.info(f"准备第 {i+1}/{len(batch_data)} 个请求")
                
                # 检查医生问题是否为空
                if not data.get("doctor_question") or not data["doctor_question"].strip():
                    logger.warning("医生问题为空，将使用默认问题")
                    data["doctor_question"] = "请描述你的症状和感受"
                
                # 解析人格信息
                target_personality = self._parse_personality_info(data["人格"])
                
                # 解析患者基本信息
                patient_info_dict = self._parse_patient_info(data["基础信息"])
                
                # 生成详细的system内容，包含患者信息、人格特征、对话历史等
                system_content = "任务：请扮演一个具有特定人格特征的患者，根据对话历史和患者信息，与医生进行门诊对话。\n"
                system_content += "重要：请严格遵守患者的人格特征。\n"
                
                # 添加患者基本信息
                if patient_info_dict:
                    system_content += "\n患者基本信息：\n"
                    for key, value in patient_info_dict.items():
                        system_content += f"- {key}：{value}\n"
                
                # 添加门诊病历信息
                if data.get("门诊病历"):
                    medical_record = data["门诊病历"]
                    system_content += "\n门诊病历：\n"
                    if isinstance(medical_record, dict):
                        for key, value in medical_record.items():
                            system_content += f"- {key}：{value}\n"
                    else:
                        system_content += f"{medical_record}\n"
                
                # 添加人格定义描述
                if target_personality:
                    personality_desc = self.prompt_generator.get_personality_description(target_personality)
                    system_content += f"\n{personality_desc}\n"
                
                # 添加对话历史
                if data["context_history"]:
                    system_content += "\n对话历史：\n"
                    for i, turn in enumerate(data["context_history"]):
                        system_content += f"轮次 {i+1}：\n"
                        if turn.get("医生"):
                            system_content += f"  医生：{turn['医生']}\n"
                        if turn.get("患者"):
                            system_content += f"  患者：{turn['患者']}\n"
                
                # 添加生成要求
                system_content += "\n\n生成要求：\n"
                system_content += "- 扮演要求：只扮演患者，回复医生。禁止回复不符合患者的任何回答。\n"
                system_content += "- 信息准确：如果回答中涉及患者信息，必须与患者基本信息中的事实性医疗信息保持一致\n"
                system_content += "- 人格体现：在回答中自然体现目标人格特征\n"
                system_content += "- 口语化表达：使用自然的口语化表达，避免过于书面化\n"
                system_content += "- 回答长度：保持合理长度，符合人格特征\n"
                system_content += "- 【输出要求】：仅生成患者的回应内容，绝对不包含本system指令中的任何内容（包括但不限于：任务描述、患者信息、人格定义、对话历史、生成要求等）\n"
                system_content += "- 【格式要求】：直接输出患者的纯对话内容，不要添加任何前缀或标识（如'患者：'等）\n"
                
                # 构建消息列表
                messages = []
                
                # 添加system消息
                messages.append({"role": "system", "content": system_content})
                
                # 添加当前医生问题作为user消息
                messages.append({"role": "user", "content": data["doctor_question"]})
                
                # 创建推理请求
                infer_request = InferRequest(messages=messages)
                infer_requests.append(infer_request)
                
                # 打印完整提示词用于检查
                logger.info(f"第 {i+1} 个请求的提示词: {messages}")
            
            # 执行批量推理
            logger.info(f"开始执行批量推理，请求数量: {len(infer_requests)}")
            resp_list = self.engine.infer(infer_requests, request_config)
            
            # 处理推理结果
            for i, resp in enumerate(resp_list):
                try:
                    raw_response = resp.choices[0].message.content
                    raw_response = self.strip_think_tags(raw_response)
                    raw_response = raw_response.strip()
                    logger.info(f"第 {i+1} 个请求的原始生成内容: {raw_response}")
                
                    responses.append(raw_response)
                except Exception as e:
                    logger.error(f"处理第 {i+1} 个请求的结果时发生错误: {str(e)}")
                    responses.append("（生成回答时发生错误）")
                
        except Exception as e:
            logger.error(f"批量生成回答失败: {str(e)}")
            # 如果出错，返回已生成的回答，并为剩余请求生成默认回答
            remaining_count = len(batch_data) - len(responses)
            if remaining_count > 0:
                responses.extend(["（生成回答时发生错误）"] * remaining_count)
        
        # 记录处理时间
        logger.info(f"批量处理完成，总耗时: {time.time() - start_time:.2f}秒")
        return responses

